package com.tich.ai.service.impl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.tich.ai.domain.entity.ChatInfo;
import com.tich.ai.domain.entity.Message;
import com.tich.ai.service.AIService;
import com.tich.constant.AIConstant;
import com.tich.constant.ExceptionConstant;
import com.tich.exception.DataAnalyseFailException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.time.Duration;
import java.util.List;


@Slf4j
@Service
public class AIServiceImpl implements AIService {

    @Autowired
    private WebClient webClient;
    @Autowired
    private ObjectMapper objectMapper;

    /**
     * 调用大模型对话接口
     * @param messages
     * @return
     */
    public Flux<String> chat(List<Message> messages) {
        //创建对话请求体
        if(!messages.isEmpty() && !messages.get(0).getRole().equals(AIConstant.SYSTEM)){
            messages.add(0, Message.builder()
                    .role(AIConstant.SYSTEM)
                    .content(AIConstant.SYSTEM_CONTENT)
                    .build());
        }
        ChatInfo chatInfo = new ChatInfo();
        chatInfo.setMessages(messages);

        return webClient.post()
                .uri(AIConstant.REQUEST_URL)
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(chatInfo)
                .retrieve()
                .onStatus(
                        status -> status.is4xxClientError() || status.is5xxServerError(),
                        response -> response.bodyToMono(String.class)
                                .flatMap(error -> Mono.error(new DataAnalyseFailException(ExceptionConstant.AI_REQUEST_ERROR + error)))
                )
                .bodyToFlux(String.class) // 原始数据流
                .timeout(Duration.ofMinutes(5)) // 总超时控制
                .concatMap(rawChunk -> {
                    // 优先处理结束信号
                    if (AIConstant.DONE.equals(rawChunk)) {
                        return Flux.just(AIConstant.DONE).concatWith(Flux.empty());
                    }
                    return parseChunk(rawChunk)
                            .timeout(Duration.ofSeconds(10)) // 单条数据处理超时
                            .flux()
                            .onErrorResume(e ->
                                    Flux.just(AIConstant.ERROR)
                            );
                }) // 数据转换
                // 全局错误处理：发送错误标识后终止
                .onErrorResume(e ->
                        Flux.just(AIConstant.ERROR)
                )
                // 确保最终终止连接
                .concatWith(Flux.defer(Flux::<String>empty));
    }

    //解析大模型的回复
    private Mono<String> parseChunk(String rawChunk) {
        return Mono.fromCallable(() -> {
            try {
                JsonNode node = objectMapper.readTree(rawChunk);
                return node.path(AIConstant.CHOICES).get(0).path(AIConstant.DELTA).path(AIConstant.CONTENT).asText();
            } catch (Exception e) {
                throw new DataAnalyseFailException(ExceptionConstant.DATA_ANALYSE_ERROR);
            }
        });
    }

}
